fix(retrieval): do not lose metric records on preemption or error - #3583
fix(retrieval): do not lose metric records on preemption or error#3583sahel-sh wants to merge 5 commits into
Conversation
The bi-encoder recipe only closed its metric loggers on the success path, and records are buffered until close(), so a run that was preempted or raised lost every metric it had produced. - close() both loggers in the finally, so an exception cannot skip them - poll StepScheduler.sigterm_received per step; both of its iterators already stop on the flag, but nothing polled it inside the step loop, so it was only checked once per epoch - size the training buffer from ckpt_every_steps, capped at DEFAULT_BUFFER_SIZE - flush both loggers: draining the buffer only reaches the file object, whose own buffer dies with the process Signed-off-by: Sahel Sharifymoghaddam <ssharifymogh@nvidia.com> (cherry picked from commit d655b4f8008f45c47470f5c6ffc022227f8167f0)
|
/ok to test bcbd19a |
There was a problem hiding this comment.
Thanks for fixing metric durability on exceptional and preempted exits. @sahel-sh A few follow-ups:
-
Please run Ruff formatting. The current lint job reports that
nemo_automodel/recipes/retrieval/train_bi_encoder.pywould be reformatted: https://github.com/NVIDIA-NeMo/Automodel/actions/runs/32207675344/job/95933955404 -
Please add focused CPU regression coverage for both critical paths introduced here:
- a signal arriving on a scheduled checkpoint step, proving the explicit post-checkpoint poll detects it and stops the loop; and
- buffered train/validation metrics being persisted when the loop raises and cleanup runs from
finally.
-
Please revisit the checkpoint-alignment claim around
buffer_size=min(DEFAULT_BUFFER_SIZE, ckpt_every_steps). Record-count buffering is not necessarily checkpoint-aligned after resuming at a nonzero step or when checkpoints occur at epoch boundaries, so a checkpoint may be persisted while its metrics remain buffered. Either flush explicitly after successful checkpoint boundaries, or narrow/correct the guarantee and test the intended behavior.
The broader StepScheduler explicit-poll/state-query redesign can remain a follow-up; I am not asking for that refactor in this PR.
svcnvidia-nemo-ci
added
the
waiting-on-customer
|
/ok to test 634afe5 |
…n coverage Addresses review feedback on the metric-durability change. Ruff formatting applied; `ruff format --check` and `ruff check` are clean on every file touched here. CHECKPOINT ALIGNMENT. The reviewer is right, and the previous comment overclaimed. buffer_size=min(DEFAULT_BUFFER_SIZE, ckpt_every_steps) counts RECORDS, and that count only coincides with checkpoint steps in the narrow case of a run starting at step 0 with purely periodic checkpoints. It drifts as soon as training resumes at a step that is not a multiple of ckpt_every_steps -- the buffer restarts empty while the checkpoint cadence does not -- and it never lines up for the checkpoints StepScheduler.is_ckpt_step also fires on: epoch boundaries, the final step, and the sigterm path. In those cases a checkpoint reaches disk while the metrics describing the steps it covers are still buffered, which is exactly what the change set out to prevent. Rather than narrow the claim, make it true: MetricLogger grows an explicit flush(), and the recipe calls it immediately after a successful save_checkpoint. buffer_size is retained but demoted to what it actually is -- a bound on how much can be pending if the process is killed without running any cleanup -- and its comment now says so instead of promising alignment. MetricLoggerDist.flush() carries the same rank-zero guard as close(). close() and flush() share a _drain() helper so their semantics cannot diverge. REGRESSION COVERAGE, CPU only, no distributed init: test_sigterm_on_checkpoint_step_saves_then_stops_the_loop -- a signal arriving on a scheduled checkpoint step still writes the checkpoint, and the post-checkpoint poll then stops the loop rather than running further steps. Asserts the loop ran exactly the steps up to the signal, the checkpoint was taken, and no metrics were lost. test_metrics_persisted_when_the_loop_raises -- the loop raises mid-step and the buffered records still reach disk, covering the move of logger.close() into the finally. test_checkpoint_flushes_metrics_when_not_buffer_aligned -- a checkpoint at a step no record-count boundary coincides with, which is the drift case above. Also pins the ordering: the flush happens after the save, not before. The broader StepScheduler explicit-poll/state-query redesign is left as the follow-up the reviewer described. Worth noting for that work: sigterm_received polls at most once per step (guarded by _sig_polled_step), so a signal that arrives DURING a long checkpoint save is not observed until the next step. The sticky flag means it is never lost, only deferred by one step, and the tests above cover the signal-before-the-step case rather than asserting the during-the-save case works. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
@yuhezhang-ai thanks for the review, and apologies for my delay: 1 and 2 are done, for 3, I made sure flush happens at successful checkpoint boundaries per your suggestion |
yuhezhang-ai
left a comment
yuhezhang-ai
left a comment
There was a problem hiding this comment.
Thanks for the thoughtful follow-up and for addressing the earlier feedback. Ruff is clean now, and the three focused CPU tests pass locally. I found one blocking runtime issue in the new flush path, plus a test gap that currently hides it:
-
MetricLogger.__init__already assignsself.flush = flush, so the newflush()method is shadowed by that per-instance boolean. On the current head,type(logger.flush)isbool, and calling it raisesTypeError: 'bool' object is not callable. Consequently, the first checkpoint will fail atself.metric_logger_train.flush(). Could you please rename the stored flag (for example,_fsync_on_write) or give the buffer operation a distinct name, and add a small test using the realMetricLogger? -
test_checkpoint_flushes_metrics_when_not_buffer_aligneddoes not currently prove the checkpoint flush occurred:at_saveis sampled before the flush, while the final persisted count is satisfied byclose()in thefinally. The test would still pass if the recipe's explicit checkpoint flush calls were removed. Please observe the real file after the checkpoint flush but before final close, or otherwise assert the intermediate durable state. -
The comment saying that flushing after checkpoint save means a checkpoint is “never on disk without metrics” is stronger than the implementation guarantees: there is still a window between save completion and the subsequent flush. Please either flush before checkpoint publication if that strict invariant is intended, or narrow the wording to the actual post-save boundary guarantee.
Finally, could you please rebase the branch onto current main with DCO sign-off on every resulting commit? The current head is bd4f63b7, and its DCO check is ACTION_REQUIRED because the latest commit has no Signed-off-by trailer.
Thank you again—the overall direction looks good once these points are addressed.
5 participants
What does this PR do ?
The bi-encoder recipe only closed its metric loggers on the success path, and records are buffered until close(), so a run that was preempted or raised lost every metric it had produced.
Changelog
Before your PR is "Ready for review"
Pre checks:
If you haven't finished some of the above items you can still open "Draft" PR.
Additional Information